#!/usr/bin/python
import numpy
import pylab
import util
import os
import math
from scipy import signal

def getz(l, i):
    try:
        return l[i]
    except:
        return 0 

def mylfilter(b, a, data):
    """
    a[0]y[n]+a[1]y(n-1)+ ... + a[N]y(n-N) = b[0]x(n) + b[1]x(n-1) + ... + b[n]x(n-N)
    """
    new_data = [0.0]*(len(data))
    for n in xrange(len(data)):
        for k in xrange(len(b)):
            new_data[n] += b[k]*getz(data, n-k)
        for k in xrange(1, len(a)):
            new_data[n] -= a[k]*getz(new_data, n-k) 
        new_data[n] = new_data[n]/a[0]
    return new_data

def src(iq_data, u, d, debug=False):
    """
        iq_data: the data
        u: upsample factor
        d: downsample factor
    """
    if debug:pylab.plot( iq_data) 
    if debug:pylab.figure()
    if debug:pylab.plot(abs(pylab.fft(iq_data)))
    
    shifted = signal.signaltools.lfilter([2, 1], [3], iq_data)
    if debug:pylab.figure()
    if debug:pylab.plot(shifted) 
    if debug:pylab.figure()
    if debug:pylab.plot(abs(pylab.fft(shifted)))

    upsample_data = []
    for i in iq_data:
        upsample_data.append(i)
        upsample_data.extend([0]*(u-1))
    if debug:pylab.figure()
    if debug:pylab.plot(upsample_data) 
    if debug:pylab.figure()
    if debug:pylab.plot(abs(pylab.fft(upsample_data)))

    wp = 0.5/u - 0.01
    ws = 0.5/u + 0.01
    gpass = 1
    gstop = 60
    N,Wn = signal.filter_design.cheb2ord(wp, ws,gpass, gstop) 
    b,a = signal.filter_design.cheby2(N, gstop, Wn)
    #if debug:print "cheby2 b %r, a %r"%(b, a)
    upsample_data = signal.signaltools.lfilter(b,a,upsample_data)
    #upsample_data = mylfilter(b,a,upsample_data)
    #b = signal.firwin(160, 1.0/u)
    #upsample_data = mylfilter(b,[1],upsample_data)
    """
    if debug:
        #signal.freqz(b, a, plot=pylab.plot)
        h,w= signal.freqz(b, a)
        fig = pylab.figure()
        pylab.title('IIR frequency response')
        ax1 = fig.add_subplot(111)
        pylab.semilogy(h,numpy.abs(w),'b')
        pylab.ylabel('Amplitude (dB)', color='b')
        pylab.xlabel('Frequency (rad/sample)')
        pylab.grid()
        pylab.legend()
        ax2=ax1.twinx()
        angles=numpy.unwrap(numpy.angle(w))
        pylab.plot(h, angles,'g')
        pylab.ylabel('Angle (radians)',color='g')

    b = signal.firwin(16, 1.0/u)
    if debug:
        h,w= signal.freqz(b)
        fig = pylab.figure()
        pylab.title('FIR frequency response')
        ax1 = fig.add_subplot(111)
        pylab.semilogy(h,numpy.abs(w),'b')
        pylab.ylabel('Amplitude (dB)', color='b')
        pylab.xlabel('Frequency (rad/sample)')
        pylab.grid()
        pylab.legend()
        ax2=ax1.twinx()
        angles=numpy.unwrap(numpy.angle(w))
        pylab.plot(h, angles,'g')
        pylab.ylabel('Angle (radians)',color='g')
    """

    if debug:pylab.figure()
    if debug:pylab.plot(upsample_data) 
    if debug:pylab.figure()
    if debug:pylab.plot(abs(pylab.fft(upsample_data)))

    downsample_data = []
    for i in xrange(len(upsample_data)):
        if 0==(i%(d-1)):
            downsample_data.append(upsample_data[i])

    if debug:pylab.figure()
    if debug:pylab.plot(downsample_data) 
    if debug:pylab.figure()
    if debug:pylab.plot(abs(pylab.fft(downsample_data)))
    if debug:pylab.show()
    return downsample_data

if __name__=='__main__':
    
    f0=1000.0
    fs = 23000.0
    
    data = [math.sin(f0*2*math.pi*i/fs) for i in xrange(190)]
    src(data, 14, 50, debug=True)
